package com.example.springbootlearn.netty.baseasr;

import cn.hutool.json.JSONUtil;
import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import okhttp3.WebSocket;
import okio.ByteString;

import java.util.Arrays;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.ScheduledFuture;


@Slf4j
@Data
public class WebSocketInfo {

    public final static String END = "{\"type\": \"end\"}";

    private WebSocket webSocket;
    private byte[] data;
    private Boolean firstMessage;
    private long lastSendTime;
    // 累计包的次数
    private int cumulativeCount;
    // 累计包的大小
    private int cumulativeSize;
    // 累计包的大小到达次数就发生一次发送
    private int needCount;

    private BlockingQueue<byte[]> queue;
    private ScheduledFuture<?> scheduledFuture;
    private CloseStrategy closeStrategy;


    public WebSocketInfo(WebSocket webSocket, int needCount, ScheduledFuture<?> scheduledFuture, CloseStrategy closeStrategy) {
        this.webSocket = webSocket;
        if (scheduledFuture == null) {
            this.needCount = -1;
        } else {
            this.needCount = needCount;
        }
        this.scheduledFuture = scheduledFuture;
        this.data = new byte[1028];
        this.lastSendTime = System.currentTimeMillis();
        this.cumulativeCount = 0;
        this.cumulativeSize = 0;
        this.queue = new LinkedBlockingQueue<>();
        this.closeStrategy = closeStrategy;
    }

    public void addQueue(byte[] packet) {
        queue.offer(packet);
    }

    public void resetPackageInfo() {
        this.lastSendTime = System.currentTimeMillis();
        this.cumulativeCount = 0;
        this.cumulativeSize = 0;
        Arrays.fill(data, (byte) 0);
    }

    public void updateLastActiveTime() {
        this.lastSendTime = System.currentTimeMillis();
    }


    public byte[] poll() {
        return queue.poll();
    }

    public void sendDataByArray(byte[] data) {
        updateLastActiveTime();
        if (this.needCount == -1) {
            sendDataDirect(data);
            return;
        }
        int length = data.length;
        log.info("send data length: " + length);
        System.arraycopy(data, 0, this.getData(), this.getCumulativeSize(), length);
        this.cumulativeSize += length;
        if (++this.cumulativeSize >= this.needCount) {
            byte[] sessionDataData = this.getData();
            byte[] sessionDataDataCopy = Arrays.copyOf(sessionDataData, sessionDataData.length);
            this.addQueue(sessionDataDataCopy);
            this.resetPackageInfo();
        }
    }

    public void sendDataDirect(byte[] data) {
        this.webSocket.send(ByteString.of(data));
    }

    public void close() {
        closeStrategy.close(this);
        if (scheduledFuture != null && !scheduledFuture.isCancelled()) {
            scheduledFuture.cancel(true);
        }
        if (queue != null && !queue.isEmpty()) {
            queue.clear();
        }
    }

    public void sendHeartBeat() {
        log.info("send heart beat");
    }

    public interface CloseStrategy {
        void close(WebSocketInfo webSocketInfo);
    }

    public static class EndMessageCloseStrategy implements CloseStrategy {
        @Override
        public void close(WebSocketInfo webSocketInfo) {
            webSocketInfo.getWebSocket().send(END);
        }
    }

    public static class StatusCodeCloseStrategy implements CloseStrategy {

        @Override
        public void close(WebSocketInfo webSocketInfo) {
            Map<String, String> header = new HashMap<>();
            header.put("namespace", "SpeechTranscriber");
            header.put("name", "StopTranscription");

            // 创建一个HashMap来表示整个JSON对象
            Map<String, Map<String, String>> jsonObject = new HashMap<>();
            jsonObject.put("header", header);
            webSocketInfo.getWebSocket().send(JSONUtil.toJsonStr(jsonObject));
        }
    }

}
